import torch
import numpy as np
import matplotlib.pyplot as plt

#随机种子，确保每次运行结果一致
torch.manual_seed(42)
X = torch.randn(100,2)#100个样本，每个样本2个特征
true_w = torch.tensor([2.0,3.0])
true_b = 4.0
Y = X @ true_w + true_b + torch.randn(100)*0.1
#打印部分数据
print(X[:5])
print(Y[:5])

import torch.nn as nn

class LinearRegressionModel(nn.Module):
    def __init__(self):
        super(LinearRegressionModel,self).__init__()
        #定义一个线性层，输入为2个特征，输出为1个预测值
        self.linear=nn.Linear(2,1)#输入维度2

model=LinearRegressionModel()

criterion=nn.MSELoss()
optimizer =torch.optim.SGD(model.parameters(),lr=0.001)

#训练模型
num_epochs=1000
for epech in range(num_epochs):
    model.train()#设置模型为训练模式
    predictions=model(x)
    loss=criterion(predictions.squeeze(),Y)#计算损失 （注意预测值需要压缩为1D）
    optimizer.zero_grad()#清空之前的梯度
    loss.backward()#计算梯度
    optimizer.step()#更新模型参数
    #打印损失
    if (epoch+1)%100==0:
        print(f'Epoch [{epoch+1}/1000],loss:{loss.item():.4f}')
"""
prodictions.squeeze() :将模型的输出从2D张量压缩为1D，目标值Y是一个一维的数组。
opimizer.zero_grad():每次反向传播前需要清空之前的梯度
loss.backward()：计算梯度。
optimizer.step():更新权重和偏执
"""

#评估模型
print(f'Predicted weight:{model.linear.weight.data.numpy()}')
print(f'Predicted bias:{model.linear.bias.data.numpy()}')
with torch.no_grad():
    predictions=model(x)
plt.scatter(X[:,0],Y,color='blue',label='True values')
plt.scatter(X[:,0],predictions,color='red',label='Predictions')
plt.legend()
plt.show()
